#include <opencv2/opencv.hpp>
#include <opencv2/dnn/dnn.hpp>
#include <iostream>
// #include <direct.h>
using namespace std;
using namespace cv;

class dnnfacedetect
{
private:
    string _modelbinary, _modeldesc;
    dnn::Net _net;

public:
    //构造函数 传入模型文件
    dnnfacedetect();
    dnnfacedetect(const string &modelBinary, const string &modelDesc);

    ~dnnfacedetect();
    //置信阈值
    float confidenceThreshold;
    double inScaleFactor;
    size_t inWidth;
    size_t inHeight;
    Scalar meanVal;

    //初始化DNN网络
    bool initdnnNet();

    //人脸检测
    vector<Mat> detect(const Mat &frame);
};

dnnfacedetect::dnnfacedetect()
{
    dnnfacedetect("", "");
}

//构造函数
dnnfacedetect::dnnfacedetect(const string &modelBinary, const string &modelDesc)
{
    _modelbinary = modelBinary;
    _modeldesc = modelDesc;

    //初始化置信阈值
    confidenceThreshold = 0.6;
    inScaleFactor = 0.5;
    inWidth = 300;
    inHeight = 300;
    meanVal = Scalar(104.0, 177.0, 123.0);
}

dnnfacedetect::~dnnfacedetect()
{
    _net.~Net();
}

//初始化dnnnet
bool dnnfacedetect::initdnnNet()
{
    _net = dnn::readNetFromTensorflow(_modelbinary, _modeldesc);
    _net.setPreferableBackend(dnn::DNN_BACKEND_OPENCV);
    _net.setPreferableTarget(dnn::DNN_TARGET_CPU);

    return !_net.empty();
}

//人脸检测
vector<Mat> dnnfacedetect::detect(const Mat &frame)
{
    Mat tmpsrc = frame.clone();
    vector<Mat> dsts;
    // 修改通道数
    if (tmpsrc.channels() == 4)
        cvtColor(tmpsrc, tmpsrc, COLOR_BGRA2BGR);
    // 输入数据调整
    Mat inputBlob = dnn::blobFromImage(tmpsrc, inScaleFactor,
                                       Size(inWidth, inHeight), meanVal, false, false);
    _net.setInput(inputBlob, "data");

    double t = (double)getTickCount();
    //人脸检测
    Mat detection = _net.forward("detection_out");

    Mat detectionMat(detection.size[2], detection.size[3],
                     CV_32F, detection.ptr<float>());

    t = (double)getTickCount() - t;
    printf("detection time = %g ms\n", t * 1000 / getTickFrequency());
    //检测出的结果进行绘制和存放到dsts中
    for (int i = 0; i < detectionMat.rows; i++)
    {
        //置值度获取
        float confidence = detectionMat.at<float>(i, 2);
        //如果大于阈值说明检测到人脸
        if (confidence > confidenceThreshold)
        {
            //计算矩形
            int xLeftBottom = static_cast<int>(detectionMat.at<float>(i, 3) * tmpsrc.cols);
            int yLeftBottom = static_cast<int>(detectionMat.at<float>(i, 4) * tmpsrc.rows);
            int xRightTop = static_cast<int>(detectionMat.at<float>(i, 5) * tmpsrc.cols);
            int yRightTop = static_cast<int>(detectionMat.at<float>(i, 6) * tmpsrc.rows);
            //生成矩形
            Rect rect((int)xLeftBottom, (int)yLeftBottom,
                      (int)(xRightTop - xLeftBottom),
                      (int)(yRightTop - yLeftBottom));

            //截出图矩形存放到dsts数组中
            Mat tmp = tmpsrc(rect);
            dsts.push_back(tmp);

            //在原图上用红框画出矩形
            rectangle(frame, rect, Scalar(0, 0, 255));
        }
    }

    return dsts;
}

int main(int argc, char **argv)
{
    //获取程序目录

    const string filepath("/data/licc/build/opencv-4.2.0/samples/dnn/face_detector");
    cout << filepath << endl;
    //定义模型文件
    string ModelBinary = filepath + "/opencv_face_detector_uint8.pb";
    string ModelDesc = filepath + "/opencv_face_detector.pbtxt";

    //图片文件
    string picdesc = "/home/licc/picture/MaYun/my2.jpg";

    cout << ModelBinary << endl;
    cout << ModelDesc << endl;

    //加载图片
    Mat frame = imread(picdesc);
    imshow("src", frame);

    try
    {
        //初始化
        dnnfacedetect fdetect = dnnfacedetect(ModelBinary, ModelDesc);
        if (!fdetect.initdnnNet())
        {
            cout << "初始化DNN人脸检测失败！" << endl;
            return -1;
        }

        if (!frame.empty())
        {
            vector<Mat> dst = fdetect.detect(frame);
            
            if (!dst.empty())
            {
                for (int i = 0; i < dst.size(); i++)
                {
                    string title = "dst" + i;
                    imshow(title, dst[i]);
                }
                imshow("src2", frame);
            }
        }
    }
    catch (const std::exception &ex)
    {
        cout << ex.what() << endl;
    }

    waitKey(0);
    return 0;
}